package com.hzgj.bcl.soa.interceptor;

import com.hzgj.bcl.soa.Event;
import com.hzgj.bcl.soa.Interceptor;
import com.hzgj.bcl.soa.rpc.message.Message;
import com.hzgj.bcl.soa.util.Bytes;
import com.hzgj.bcl.soa.util.Constants;
import com.hzgj.bcl.util.jmx.annotation.MBean;
import com.hzgj.bcl.util.jmx.annotation.ManagedAttribute;

import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.zip.DataFormatException;
import java.util.zip.Deflater;
import java.util.zip.Inflater;

import lombok.Setter;

/**
 * Compresses the payload of a message. Goal is to reduce the number of messages
 * sent across the wire. Should ideally be layered somewhere below serialization.
 */
@MBean(description = "Compresses messages to hanle and uncompresses received messages", objectName = Constants.OBJECTNAME_PREFIX + Constants.OBJECTNAME_INTERCEPTOR + "," +
        "name=Compress")
public class Compress extends Interceptor {

    public static final String MESSAGE_HEADER_COMPRESS = "message.header.compress";

    /* -----------------------------------------    Properties     -------------------------------------------------- */

    @ManagedAttribute(description = "Compression level (from java.util.zip.Deflater) " +
            "(0=no compression, 1=best speed, 9=best compression). Default is 9")
    protected int compression_level = Deflater.BEST_COMPRESSION; // this is 9

    @ManagedAttribute(description = "Minimal payload size of a message (in bytes) for compression to kick in. Default is 500 bytes")
    @Setter
    protected long min_size = 500;

    @ManagedAttribute(description = "Number of inflaters/deflaters for concurrent processing. Default is 2 ")
    protected int pool_size = 2;


    protected BlockingQueue<Deflater> deflater_pool = null;
    protected BlockingQueue<Inflater> inflater_pool = null;


    @Override
    public void start() throws Exception {
        deflater_pool = new ArrayBlockingQueue<>(pool_size);
        for (int i = 0; i < pool_size; i++) {
            deflater_pool.add(new Deflater(compression_level));
        }
        inflater_pool = new ArrayBlockingQueue<>(pool_size);
        for (int i = 0; i < pool_size; i++) {
            inflater_pool.add(new Inflater());
        }
    }

    @Override
    public void destroy() {
        for (Deflater deflater : deflater_pool) {
            deflater.end();
        }
        for (Inflater inflater : inflater_pool) {
            inflater.end();
        }
    }


    /**
     * We compress the payload if it is larger than <code>min_size</code>. In this case we add a header containing
     * the original size before compression. Otherwise we add no header.<br/>
     * Note that we compress either the entire buffer (if offset/length are not used), or a subset (if offset/length
     * are used)
     *
     * @param evt
     */
    @Override
    public Object down(Event evt) {
        if (evt.getType() == Event.MSG) {
            Message msg = (Message) evt.getArg();
            int length = msg.getLength(); // takes offset/length (if set) into account
            if (length >= min_size) {
                byte[] payload = (byte[]) msg.getPayload(); // here we get the ref so we can avoid copying
                byte[] compressed_payload = new byte[length];
                Deflater deflater = null;
                try {
                    deflater = deflater_pool.take();
                    deflater.reset();
                    deflater.setInput(payload, 0, length);
                    deflater.finish();
                    deflater.deflate(compressed_payload);
                    int compressed_size = deflater.getTotalOut();

                    if (compressed_size < length) {
                        byte[] new_payload = new byte[compressed_size + 4];//包含原始字节的长度
                        System.arraycopy(compressed_payload, 0, new_payload, 4, compressed_size);
                        System.arraycopy(Bytes.int2bytes(length), 0, new_payload, 0, 4);
                        Message copy = msg.copy(false).setPayload(new_payload).putHeader(MESSAGE_HEADER_COMPRESS, (byte) 1);
                        if (log.isTraceEnabled()) {
                            log.trace("down(): compressed payload from " + length + " bytes to " + compressed_size + " bytes");
                        }
                        return down_int.down(new Event(Event.MSG, copy));
                    } else {
                        if (log.isTraceEnabled()) {
                            log.trace("down(): skipping compression since the compressed message (" + compressed_size +
                                    ") is not smaller than the original (" + length + ")");
                        }
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt(); // set interrupt flag again
                    throw new RuntimeException(e);
                } finally {
                    if (deflater != null) {
                        deflater_pool.offer(deflater);
                    }
                }
            }
        }
        return down_int.down(evt);
    }


    /**
     * If there is no header, we pass the message up. Otherwise we uncompress the payload to its original size.
     *
     * @param evt
     */
    @Override
    public Object up(Event evt) {
        if (evt.getType() == Event.MSG) {
            Message msg = (Message) evt.getArg();
            Byte compress = (Byte) msg.getHeader(MESSAGE_HEADER_COMPRESS);
            if (compress != null && compress == (byte) 1) {
                Message uncompressed_msg = uncompress(msg);
                if (uncompressed_msg != null) {
                    if (log.isTraceEnabled()) {
                        log.trace("up(): uncompressed " + msg.getLength() + " bytes to " + uncompressed_msg.getLength() + " bytes");
                    }
                    return up_int.up(new Event(Event.MSG, uncompressed_msg));
                }
            }
        }
        return up_int.up(evt);
    }

    /**
     * Returns a new message as a result of uncompressing msg, or null if msg couldn't be uncompressed
     */
    protected Message uncompress(Message msg) {
        byte[] compressed_payload = (byte[]) msg.getPayload();
        if (compressed_payload != null && compressed_payload.length > 0) {
            byte[] length = new byte[4];
            System.arraycopy(compressed_payload, 0, length, 0, 4);
            int original_size = Bytes.bytes2int(length);
            byte[] uncompressed_payload = new byte[original_size];
            Inflater inflater = null;
            try {
                inflater = inflater_pool.take();
                inflater.reset();
                inflater.setInput(compressed_payload, 4, msg.getLength());
                try {
                    inflater.inflate(uncompressed_payload);
                    return msg.copy(false).setPayload(uncompressed_payload);
                } catch (DataFormatException e) {
                    log.error("{}", e);
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt(); // set the interrupt bit again, so caller can handle it
            } finally {
                if (inflater != null) {
                    inflater_pool.offer(inflater);
                }
            }
        }
        return null;
    }
}
